# -*- coding: utf-8 -*-
#
#    Powered by AIgames of michine intellengence lab
#    Author: Justin Xu
# =================================================================
"""
    Read the go's records from recods.file
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

from six.moves import xrange  # pylint: disable=redefined-builtin
import tensorflow as tf

import psigma_go
IMAGE_SIZE = 19

NUM_CLASSES =361
NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN = 2000000
NUM_EXAMPLES_PER_EPOCH_FOR_EVAL = 100000

TRAIN_FILE = 'train.tfrecords'
VALIDATION_FILE = 'validation.tfrecords'

def read_and_decode(filename_queue):
  reader = tf.TFRecordReader()
  _, serialized_example = reader.read(filename_queue)
  features = tf.parse_single_example(
      serialized_example,
      # Defaults are not specified since both keys are required.
      features={
          'image_raw': tf.FixedLenFeature([], tf.string),
          'label': tf.FixedLenFeature([], tf.int64),
      })

  # Convert from a scalar string tensor
  image = tf.decode_raw(features['image_raw'], tf.int8)
  # import pdb; pdb.set_trace()
  image.set_shape([psigma_go.IMAGE_PIXELS])
  image = tf.reshape(image, [19,19,2], name=None)
  # OPTIONAL: Could reshape into a 19x19 image and apply distortions
  # here.  Since we are not applying any distortions in this
  # example, and the next step expects the image to be flattened
  # into a vector, we don't bother.

  # Convert from [0, 255] -> [-0.5, 0.5] floats.
  image = tf.cast(image, tf.float32)

  # Convert label from a scalar uint8 tensor to an int32 scalar.
  # import pdb; pdb.set_trace()
  label = tf.cast(features['label'], tf.int32)
  return image, label

def distorted_inputs(batch_size):
  """Reads input data num_epochs times.

  Args:
    train: Selects between the training (True) and validation (False) data.
    batch_size: Number of examples per returned batch.
    num_epochs: Number of times to read the input data, or 0/None to
       train forever.

  Returns:
    A tuple (images, labels), where:
    * images is a float tensor with shape [batch_size, []]
      in the range [-0.5, 0.5].
    * labels is an int32 tensor with shape [batch_size] with the true label,
      a number in the range [0, mnist.NUM_CLASSES).
    Note that an tf.train.QueueRunner is added to the graph, which
    must be run using e.g. tf.train.start_queue_runners().
  """

  # if not num_epochs: num_epochs = None
  filename = os.path.join('./',TRAIN_FILE)
  filename_queue = tf.train.string_input_producer([filename])
    # Even when reading in multiple threads, share the filename
    # queue.
  image, label = read_and_decode(filename_queue)
    # import pdb; pdb.set_trace()
    # image = image.reshape((19,19,2))
  image.set_shape([19,19,2])
  # import pdb; pdb.set_trace()
    # label.set_shape([1])

    # Ensure that the random shuffling has good mixing properties
  min_fraction_of_examples_in_queue = 0.07
  min_queue_examples = int(NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN *
                             min_fraction_of_examples_in_queue)

    # Shuffle the examples and collect them into batch_size batches.
    # (Internally uses a RandomShuffleQueue.)
    # We run this in two threads to avoid being a bottleneck.
  images, sparse_labels = tf.train.shuffle_batch(
        [image, label], batch_size=batch_size, num_threads=24,
        capacity=1000 + 2 * batch_size,
        # Ensures a minimum amount of shuffling of examples.
        min_after_dequeue=1000)

    # Display the training images in the visualizer
  # tf.summary.image('images', images)
  return images, tf.reshape(sparse_labels, [batch_size])

def inputs(batch_size):
  """Reads input data num_epochs times.

  Args:
    train: Selects between the training (True) and validation (False) data.
    batch_size: Number of examples per returned batch.
    num_epochs: Number of times to read the input data, or 0/None to
       train forever.

  Returns:
    A tuple (images, labels), where:
    * images is a float tensor with shape [batch_size, []]
      in the range [-0.5, 0.5].
    * labels is an int32 tensor with shape [batch_size] with the true label,
      a number in the range [0, mnist.NUM_CLASSES).
    Note that an tf.train.QueueRunner is added to the graph, which
    must be run using e.g. tf.train.start_queue_runners().
  """

  #  if not num_epochs: num_epochs = None
  filename = os.path.join('./',VALIDATION_FILE)
  filename_queue = tf.train.string_input_producer([filename])
    # Even when reading in multiple threads, share the filename
    # queue.
  image, label = read_and_decode(filename_queue)
    # import pdb; pdb.set_trace()
    # image = image.reshape((19,19,2))
  image.set_shape([19,19,2])
  # import pdb; pdb.set_trace()
  # label.set_shape([1])

  # Ensure that the random shuffling has good mixing properties
  min_fraction_of_examples_in_queue = 0.07
  min_queue_examples = int(NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN *
                           min_fraction_of_examples_in_queue)

  # Shuffle the examples and collect them into batch_size batches.
  # (Internally uses a RandomShuffleQueue.)
  # We run this in two threads to avoid being a bottleneck.
  images, sparse_labels = tf.train.shuffle_batch(
      [image, label], batch_size=batch_size, num_threads=24,
      capacity=1000 + 2 * batch_size,
      # Ensures a minimum amount of shuffling of examples.
      min_after_dequeue=1000)

  # Display the training images in the visualizer
  # tf.summary.image('images', images)
  return images, tf.reshape(sparse_labels, [batch_size])
